# uav.py
import numpy as np

class UAV:
    def __init__(self, energy_states, cruise_states, num_actions,charging_stations ,init_energy, init_position,id, learning_rate=0.1, gamma=0.9, epsilon=0.1, grid_size=1000,speed = 25):
        # 初始化二维Q表，行表示状态，列表示动作
        self.energy_Q = np.zeros((energy_states, 2))  # 能量策略Q表
        self.cruise_Q = np.zeros((cruise_states, num_actions))  # 飞行策略Q表
        self.learning_rate = learning_rate
        self.gamma = gamma
        self.epsilon = epsilon
        self.speed = speed
        self.id = id
        self.has_charged = False
        
        # 无人机的初始电量和位置
        self.energy = init_energy  # 当前电量
        self.position = init_position  # 当前位置，表示为(x, y)的元组
        self.grid_size = grid_size  # 网格大小，限制移动范围
        self.charging_stations = charging_stations
        
        # 历史路径记录
        self.history_position = []
        self.history_length = 10
        
    
    def choose_action(self, state, Q_table):
        # Epsilon-greedy策略选择动作
        if np.random.uniform(0, 1) < self.epsilon:
            return np.random.choice(Q_table.shape[1])  # 随机选择一个动作
        return np.argmax(Q_table[state])  # 选择当前状态下的最佳动作

    def update_Q(self, Q_table, state, action, reward, next_state):
        # 确保状态和动作在Q表范围内
        if state < Q_table.shape[0] and action < Q_table.shape[1]:
            # 限制next_state的值在合法范围内
            next_state = max(0, min(next_state, Q_table.shape[0] - 1))

            # Q-learning更新规则
            predict = Q_table[state][action]
            target = reward + self.gamma * np.max(Q_table[next_state])
            Q_table[state][action] += self.learning_rate * (target - predict)

    def select_energy_action(self, energy_state):
        # 根据能量状态选择动作，0表示充电，1表示不充电
        return self.choose_action(energy_state, self.energy_Q)

    def select_cruise_action(self, cruise_state,other_uavs_positions, uav_position):
        # 根据飞行状态选择动作
        # 获取所有可能的动作
        possible_actions = [0, 1, 2, 3]  # 假设0-3为移动方向
        valid_actions = []
        x,y = uav_position

        # 检查每个可能动作是否会导致冲突
        for action in possible_actions:
            new_position = self.predict_position(action)  # 预测该动作会导致的新位置
            flag = 0
            if new_position not in other_uavs_positions and new_position not in self.history_position:
                if y<=25 and action ==1:
                    flag =1
                if  y >= 975  and action == 0:
                    flag = 1
                if x <= 25 and action ==2:
                    flag = 1
                if x >= 975 and action ==3:
                    flag = 1 
            if flag == 0:    
                valid_actions.append(action)  # 无冲突，将动作加入有效动作列表

        # 从有效动作中选择最佳动作
        if np.random.uniform(0, 1) < self.epsilon:
            return np.random.choice(valid_actions)  # 随机选择一个有效动作
        else:
            # 根据 Q 表选择最佳有效动作
            best_action = valid_actions[0]
            max_q_value = self.cruise_Q[cruise_state][best_action]
            for action in valid_actions:
                if self.cruise_Q[cruise_state][action] > max_q_value:
                    best_action = action
                    max_q_value = self.cruise_Q[cruise_state][action]
        
        return best_action

    def predict_position(self, action):
            """
            根据动作预测新位置，而不更新当前的位置
            """
            x, y = self.position

            # 根据动作预测新位置
            if action == 0:  # 向上移动
                y += 1
            elif action == 1:  # 向下移动
                y -= 1
            elif action == 2:  # 向左移动
                x -= 1
            elif action == 3:  # 向右移动
                x += 1
            # 停留位置不变

            # 限制边界
            x = max(0, min(x, self.grid_size - 1))
            y = max(0, min(y, self.grid_size - 1))
            
            return (x, y)
    def update_energy_Q(self, energy_state, action, reward, next_energy_state):
        # 更新能量策略Q表
        self.update_Q(self.energy_Q, energy_state, action, reward, next_energy_state)

    def update_cruise_Q(self, cruise_state, action, reward, next_cruise_state):
        # 更新飞行策略Q表
        self.update_Q(self.cruise_Q, cruise_state, action, reward, next_cruise_state)

    def execute_cruise_action(self, action):
        """
        根据动作更新位置并消耗对应电量
        假设动作0-3分别表示上下左右移动，动作4表示原地停留
        """
        # 定义每个动作的电量消耗（例如，上下左右移动消耗5单位电量，停留消耗2单位电量）
        
        # 当前坐标
        x, y = self.position
        distance_per_time = self.speed


        # 根据动作更新位置
        if action == 0 and y < self.grid_size - 1:  # 向上移动
            y += distance_per_time
        elif action == 1:  # 向下移动
            y -= distance_per_time
        elif action == 2:  # 向左移动
            x -= distance_per_time
        elif action == 3:  # 向右移动
            x += distance_per_time
        # 停留位置不变
        
        # 检查是否超出边界并限制
        if x < 0:
            x = 0  # 如果x小于0，设为最小值0
        elif x >= self.grid_size:
            x = self.grid_size - 1  # 如果x超出grid_size，设为最大值grid_size-1
        
        if y < 0:
            y = 0  # 如果y小于0，设为最小值0
        elif y >= self.grid_size:
            y = self.grid_size - 1  # 如果y超出grid_size，设为最大值grid_size-1
        
        # 更新位置
        self.position = (x, y)
        
        # 更新历史路径
        self.history_position.append(self.position)
        if len(self.history_position) > self.history_length:
            self.history_position.pop(0)
        
        # 执行电量消耗
        self.reduce_energy(1)

    def reduce_energy(self, amount):
        # 扣除电量
        self.energy -= amount
        if self.energy < 0:
            self.energy = 0  # 确保电量不为负数

    def valid_cruise_action(self, action):
        """
        根据当前位置和动作，检查动作是否有效，确保不会越界
        """
        x, y = self.position
        
        if action == 0:  # 向上
            if y + 1 >= self.grid_size:  # 超出上边界
                return False
        elif action == 1:  # 向下
            if y - 1 < 0:  # 超出下边界
                return False
        elif action == 2:  # 向左
            if x - 1 < 0:  # 超出左边界
                return False
        elif action == 3:  # 向右
            if x + 1 >= self.grid_size:  # 超出右边界
                return False
        return True

    def select_valid_cruise_action(self, cruise_state):
        # 尝试选择一个有效的动作
        valid_actions = []
        
        for action in range(4):  # 假设有4个移动方向
            if self.valid_cruise_action(action):
                valid_actions.append(action)
        
        if not valid_actions:
            return 4  # 如果没有有效的动作，选择停留动作
        
        # 从有效动作中选择一个
        return np.random.choice(valid_actions)

    def serve_ground_user(self, ground_user_position,communacation_demand):
        """
        为地面用户提供服务，计算服务的电量消耗和奖励
        """
        distance = np.linalg.norm(np.array(self.position) - np.array(ground_user_position))
        
        # 假设服务电量消耗与距离成正比
        energy_cost = distance * communacation_demand /1000   # 每单位距离消耗2单位电量
        
        if self.energy >= energy_cost:
            # 可以为用户提供服务
            self.reduce_energy(energy_cost)
            return 1  # 服务成功，奖励1
        else:
            return -1  # 电量不足，不能提供服务
        
    def find_nearest_charging_station(self):
        """
        计算无人机当前位置到每个充电站的距离，并返回最近的充电站的位置
        """
        min_distance = float('inf')
        nearest_station = None

        for station in self.charging_stations:
            distance = np.linalg.norm(np.array(self.position) - np.array(station))  # 使用欧几里得距离
            if distance < min_distance:
                min_distance = distance
                nearest_station = station
                
        uav_x ,uav_y = self.position
        station_x,station_y = nearest_station
        if abs(uav_x-station_x) + abs(uav_y-station_y) < self.speed:
            one_can_get = True
        else:
            one_can_get = False
            
        return nearest_station,one_can_get

    def charge_if_needed(self):
        """
        如果电量低于阈值，飞向最近的充电站并进行充电
        """
        # 判断电量是否低于阈值
        if self.energy < 14:  # 你可以调整这个阈值
            print(f"UAV {self.id} at {self.position} has low energy. Heading to nearest charging station.")
            self.go_to_nearest_charging_station()
            
    def go_to_nearest_charging_station(self):
        nearest_station,one_can_get = self.find_nearest_charging_station()       
        if one_can_get:
            self.position = nearest_station
        else:
            # 选择飞向最近充电站的动作
            if self.position != nearest_station:
                # 计算方向并选择一个合法的移动动作
                if nearest_station[0] > self.position[0]:  # 向右移动
                    action = 3
                elif nearest_station[0] < self.position[0]:  # 向左移动
                    action = 2
                elif nearest_station[1] > self.position[1]:  # 向上移动
                    action = 0
                else:  # 向下移动
                    action = 1
            # 执行飞行操作
            self.execute_cruise_action(action)
            
        # 保证电量为满电
        if self.position == nearest_station:
            self.energy = 100
            print(f"UAV has reached the charging station at {nearest_station} and is now fully charged.")